from datetime import datetime
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure, OperationFailure
from collections import Counter

class MongoDBManager:
    def __init__(self, url="mongodb://localhost:27017/", db_name="douyin"):
        """
        初始化MongoDBManager类。

        :param uri: MongoDB服务器的连接URI。
        :param db_name: 要连接的数据库名称。
        """
        self.url = url
        self.db_name = db_name
        self.client = None
        self.db = None
        self.connect()

    def connect(self):
        """
        连接到MongoDB数据库。
        """
        try:
            # 创建MongoDB客户端
            self.client = MongoClient(self.url)
            # 返回指定数据库
            self.db = self.client[self.db_name]
        except ConnectionFailure as e:
            print(f"连接失败: {e}")
        except Exception as e:
            print(f"发生错误: {e}")


    def count_documents(self,collection_name):
        '''获取数据库中表的计数'''
        return self.db[collection_name].count_documents({})
        
    def is_document_exists(self, collection_name, filter):
        """
        检查指定集合中是否存在与给定过滤条件匹配的文档。

        :param collection_name: 集合名称。
        :param filter: 用于筛选文档的查询条件。
        :return: 如果文档存在返回True，否则返回False。
        """
        try:
            # 检查文档是否已存在
            existing_document = self.db[collection_name].find_one(filter)
            return existing_document is not None
        except OperationFailure as e:
            print(f"检查文档存在性失败: {e}")
            return False
        
        
    def insert_one(self, collection_name, document):
        """
        向指定集合插入一个文档。

        :param collection_name: 集合名称。
        :param document: 要插入的文档。
        :return: 插入结果。
        """
        self.db[collection_name].insert_one(document)
            
            
            
    def update_one(self, collection_name, filter, update):
        """
        更新指定集合中的一个文档。如果文档不存在，则插入新文档。

        :param collection_name: 集合名称。
        :param filter: 用于筛选要更新文档的查询。
        :param update: 更新操作。
        :return: 更新结果。
        """
        try:
            result = self.db[collection_name].update_one(filter, update, upsert=True)
            if result.matched_count > 0:
                print(f"更新文档成功，匹配的文档数量: {result.matched_count}, 修改的文档数量: {result.modified_count}")
            else:
                print("没有找到匹配的文档，已插入新文档")
            return result
        except OperationFailure as e:
            print(f"更新文档失败: {e}")



    def insert_many(self, collection_name, documents):
        """
        向指定集合批量插入多个文档。

        :param collection_name: 集合名称。
        :param documents: 要插入的文档列表。
        :return: 批量插入结果。
        """
        try:
            result = self.db[collection_name].insert_many(documents)
            # print(f"批量插入文档成功，插入的文档数量: {len(result.inserted_ids)}")
            return result
        except OperationFailure as e:
            print(f"批量插入文档失败: {e}")

    def find_all(self, collection_name, attribute):
        """
        查询指定集合的所有文档，并提取指定属性的值。

        :param collection_name: 集合名称。
        :param attribute: 文档中要提取的属性名称。
        :return: 包含指定属性值的列表。
        """
        try:
            # 创建一个空列表来存储属性值
            attribute_values = []
            # 查询所有文档
            for document in self.db[collection_name].find():
                # 将属性值添加到列表中
                if attribute in document:
                    attribute_values.append(document[attribute])
                else:
                    print(f"属性 '{attribute}' 不在文档中")
            return attribute_values
        except OperationFailure as e:
            print(f"查询文档失败: {e}")
            return []
        
        
    def find_all_document(self, collection_name):
        """
        查询指定集合的所有文档。

        :param collection_name: 集合名称。
        :param attribute: 文档中要提取的属性名称。
        :return: 包含指定属性值的文档列表。
        """
        try:
            # 创建一个空列表来存储属性值
            documents = []
            # 查询所有文档
            for document in self.db[collection_name].find():
                # 将属性值添加到列表中
                documents.append(document)
            return documents
        except OperationFailure as e:
            print(f"查询文档失败: {e}")
            return []

        
        
    def find_one(self, collection_name, filter, attribute):
        """
        查询指定集合中符合条件的单个文档，并返回指定属性的值。

        :param collection_name: 集合名称。
        :param filter: 用于筛选文档的查询条件。
        :param attribute: 文档中要提取的属性名称。
        :return: 指定属性的值，如果文档不存在或属性不存在，则返回 None。
        """
        try:
            # 查询单个文档
            document = self.db[collection_name].find_one(filter)
            if document is not None and attribute in document:
                return document[attribute]
            else:
                print(f"未找到符合条件的文档或属性 '{attribute}' 不存在")
                return None
        except OperationFailure as e:
            print(f"查询文档失败: {e}")
            return None

            
    def find_filter(self, collection_name,filter):
        """
        查询指定集合中符合filter条件的文档，并返回这些文档的 `url`。

        :param collection_name: 集合名称。
        :return: 包含符合 filter 条件的文档的 `url` 的列表。
        """
        try:
            documents = self.db[collection_name].find(filter)
            # 创建一个空列表来存储符合条件的文档的 `url`
            black_true_urls = []
            # 将符合条件的文档的 `url` 添加到列表中
            for document in documents:
                # 假设 `url` 是文档中的一个属性
                url = document.get("url")
                if url:
                    black_true_urls.append(url)
            return black_true_urls
        except OperationFailure as e:
            print(f"查询文档失败: {e}")
            return []
        
        
            # 插入新的URL列表的方法
    def insert_db_urls(self, urls: list[str]):
        '''
        用于插入新的URL列表，如果URL已存在则不插入。

        参数:
            - collection_name: 集合名称
            - urls: 包含URL字典的列表
        '''
        # 过滤掉已经存在的URL
        new_urls = [{'url': url, 'black': False} for url in urls if not self.is_document_exists(
            'tk_forerver_urls', {"url": url})]

        # 插入新的URL
        if new_urls:
            self.insert_many('tk_forerver_urls', new_urls)
        return {'code': 200, 'messge': f'{new_urls} 已成功插入'}
    
        
    def get_db_black_urls(self) -> list:
        '''用于获取数据库中存储的url列表'''
        urls = self.find_filter(collection_name='tk_forerver_urls', filter = {"black": True})
        return urls
    
    def set_blacklist(self,url, collection_name):
        """
        在指定集合中，将与给定 `url` 匹配的文档的 `black` 属性设置为 `True`。

        :param url: 要设置为黑名单的 `url`。
        :param collection_name: 集合名称。
        :return: 操作结果。
        """
        try:
            filter = {"url": url}
            update = {"$set": {"black": True}}
            result = self.db[collection_name].update_one(filter, update)
            if result.matched_count > 0:
                print(f"已成功将 {url} 设置为黑名单")
            else:
                print(f"未找到与 {url} 匹配的文档")
            return result
        except OperationFailure as e:
            print(f"设置黑名单失败: {e}")
            return None
        
        
    def get_day_room_count(self) -> dict:
        current_date = datetime.now().strftime("%Y/%m/%d")
        lst = self.find_filter(collection_name='count_collection',filter={'day':current_date})
        count = Counter(lst)
        count_dict = dict(count)  # 将Counter对象转换为字典
        return count_dict
    
    def get_history_room_count(self)-> dict:
        lst = self.find_all(collection_name='count_collection',attribute='url')
        count = Counter(lst)
        count_dict = dict(count)  # 将Counter对象转换为字典
        return count_dict


    
    def close(self):
        """
        关闭MongoDB客户端连接。
        """
        if self.client is not None:
            self.client.close()
            print("数据库连接已关闭")
            
            